from django.shortcuts import render
from langchain_community.document_loaders import PyPDFLoader
import os
import base64


def pdf_load(filepath):
    # 检查文件是否存在
    if not os.path.exists(filepath):
        print(f"文件 {filepath} 不存在")
        return []

    loader = PyPDFLoader(filepath)
    documents = loader.load()

    # 遍历所有页面并打印内容
    for i, doc in enumerate(documents):
        print(f"--- Page {i + 1} ---")
        print(doc.page_content)
        print("\n")

    return documents


def extract_tables_from_pdf(filepath):
    """
    从PDF中提取表格数据
    """
    try:
        import pdfplumber
        # 检查文件是否存在
        if not os.path.exists(filepath):
            print(f"文件 {filepath} 不存在")
            return []

        tables = []
        with pdfplumber.open(filepath) as pdf:
            for page in pdf.pages:
                page_tables = page.extract_tables()
                for table in page_tables:
                    tables.append(table)
        return tables
    except ImportError:
        print("请安装pdfplumber库: pip install pdfplumber")
        return []


def extract_images_from_pdf(filepath, output_dir="extracted_images"):
    """
    从PDF中提取图片并保存到指定目录
    """
    try:
        import fitz  # PyMuPDF
        # 检查文件是否存在
        if not os.path.exists(filepath):
            print(f"文件 {filepath} 不存在")
            return []

        # 创建输出目录
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        images = []
        pdf_document = fitz.open(filepath)

        for page_num in range(len(pdf_document)):
            page = pdf_document[page_num]
            image_list = page.get_images(full=True)

            for img_index, img in enumerate(image_list):
                xref = img[0]
                base_image = pdf_document.extract_image(xref)
                image_bytes = base_image["image"]
                image_ext = base_image["ext"]

                # 生成图片文件名
                image_filename = f"page_{page_num + 1}_img_{img_index + 1}.{image_ext}"
                image_path = os.path.join(output_dir, image_filename)

                # 保存图片到文件
                with open(image_path, "wb") as image_file:
                    image_file.write(image_bytes)

                image_info = {
                    "page": page_num + 1,
                    "index": img_index + 1,
                    "ext": image_ext,
                    "filename": image_filename,
                    "path": image_path,
                    "size": len(image_bytes)
                }
                images.append(image_info)
                print(f"已保存图片: {image_path} (大小: {len(image_bytes)} 字节)")

        pdf_document.close()
        return images
    except ImportError:
        print("请安装PyMuPDF库: pip install PyMuPDF")
        return []


def display_image_info(images):
    """
    显示提取到的图片信息
    """
    if not images:
        print("没有提取到任何图片")
        return

    print(f"总共提取到 {len(images)} 张图片:")
    for img in images:
        print(f"  页面: {img['page']}, "
              f"序号: {img['index']}, "
              f"格式: {img['ext']}, "
              f"文件名: {img['filename']}, "
              f"大小: {img['size']} 字节, "
              f"路径: {img['path']}")


def get_image_base64(image_path):
    """
    将图片转换为base64编码
    """
    try:
        with open(image_path, "rb") as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
            return encoded_string
    except Exception as e:
        print(f"转换图片到base64时出错: {e}")
        return None


def find_relevant_images(query, images):
    """
    根据用户查询找到相关的图片
    这是一个简化版本，实际应用中可以使用更复杂的图像识别技术
    """
    # 简单的关键词匹配逻辑
    query_keywords = query.lower()

    # 定义关键词到图片类型的映射
    keyword_mapping = {
        "流程图": ["流程", "流程图", "步骤"],
        "表格": ["表格", "数据", "统计"],
        "截图": ["界面", "截图", "屏幕"],
        "图": ["图", "图片", "图像"],
        "照片": ["照片", "图片", "照片"],
    }

    # 查找匹配的图片
    relevant_images = []
    for keyword, synonyms in keyword_mapping.items():
        if any(synonym in query_keywords for synonym in synonyms):
            # 在已提取的图片中查找匹配的
            for img in images:
                # 这里可以添加更复杂的匹配逻辑
                relevant_images.append(img)
            break

    # 如果没有特定匹配，返回所有图片
    if not relevant_images and images:
        relevant_images = images[:3]  # 只返回前3张图片

    return relevant_images


# Create your views here.
if __name__ == '__main__':
    # 使用一个存在的PDF文件进行测试
    # 注意：你需要提供一个实际存在的PDF文件路径
    pdf_file_path = 'Boss 直聘平台使用指南.pdf'  # 请替换为实际的PDF文件路径

    if os.path.exists(pdf_file_path):
        print("开始提取PDF内容...")
        docmounts = pdf_load(pdf_file_path)
        print("文档内容提取完成")

        # 提取表格
        tables = extract_tables_from_pdf(pdf_file_path)
        print(f"提取到 {len(tables)} 个表格")
        for i, table in enumerate(tables):
            print(f"表格 {i + 1}:")
            for row in table:
                print(row)
            print("\n")

        # 提取图片
        images = extract_images_from_pdf(pdf_file_path)
        display_image_info(images)

        # 模拟用户查询
        user_query = "我需要查看相关的流程图"
        relevant_images = find_relevant_images(user_query, images)

        print(f"\n根据查询 '{user_query}' 找到的相关图片:")
        for img in relevant_images:
            print(f"  - {img['filename']} (页面 {img['page']})")
            # 获取base64编码的图片
            base64_image = get_image_base64(img['path'])
            if base64_image:
                print(f"    Base64预览: {base64_image[:50]}...")
    else:
        print(f"测试文件 {pdf_file_path} 不存在，请提供一个有效的PDF文件路径")